from timdex_dataset_api import TIMDEXDataset
from tests.utils import generate_sample_records

td = TIMDEXDataset("s3://timdex/dataset")

writes = [
    generate_sample_records(100, source='alma', run_type='full', run_date='2025-07-01', timdex_record_id_prefix='alma'),
    generate_sample_records(50, source='alma', run_date='2025-07-02', timdex_record_id_prefix='alma'),
    generate_sample_records(50, source='libguides', run_type='full', run_date='2025-07-11', timdex_record_id_prefix='libguies'),
]
for records in writes:
    td.write(records)

from timdex_dataset_api import TIMDEXDataset
from  timdex_dataset_api.config import configure_dev_logger

configure_dev_logger()

td = TIMDEXDataset("s3://timdex/dataset")
td.load(current_records=True)
"""
INFO:timdex_dataset_api.metadata:configuring S3 connection
INFO:timdex_dataset_api.metadata:creating table of full dataset metadata
INFO:timdex_dataset_api.metadata:'records' table created - rows: 200, elapsed: 0.04554949980229139
INFO:timdex_dataset_api.metadata:creating view of current records metadata
INFO:timdex_dataset_api.metadata:'current_records' view created - rows: 150, elapsed: 0.007222665939480066
INFO:timdex_dataset_api.metadata:metadata database setup elapsed: 0.09048424987122416, path: ':memory:'
INFO:timdex_dataset_api.dataset:Dataset successfully loaded: 's3://timdex/dataset', 0.11s
"""

td.metadata.conn.query("""select source, count(*) from records group by source;""")
"""
Out[3]: 
┌───────────┬──────────────┐
│  source   │ count_star() │
│  varchar  │    int64     │
├───────────┼──────────────┤
│ alma      │          150 │
│ libguides │           50 │
└───────────┴──────────────┘
"""

td.metadata.conn.query("""select source, count(*) from current_records group by source;""")
"""
Out[4]: 
┌───────────┬──────────────┐
│  source   │ count_star() │
│  varchar  │    int64     │
├───────────┼──────────────┤
│ alma      │          100 │ <------- 50 fewer records because not all current
│ libguides │           50 │
└───────────┴──────────────┘
"""